from .method_base import CenterDetectMethodBase
import numpy as np
import scipy.signal as ss


class ConvenCentroid(CenterDetectMethodBase):
    def _process(self, *args, **kwargs):
        # 生成横坐标序列
        row_axis = np.ones([self.height, 1], dtype=np.uint32) @ np.arange(
            self.axis_0, self.width + self.axis_0, dtype=np.uint32
        ).reshape([1, self.width])
        # 生成纵坐标序列
        col_axis = np.arange(
            self.axis_0, self.height + self.axis_0, dtype=np.uint32
        ).reshape([self.width, 1]) @ np.ones([1, self.width], dtype=np.uint32)
        # 估计图像中心横坐标：
        x_c = np.sum(row_axis * self.fig) / np.sum(self.fig)
        # 估计图像中心纵坐标：
        y_c = np.sum(col_axis * self.fig) / np.sum(self.fig)
        return x_c, y_c

    @property
    def _name(self):
        return "ConvenCentroid"


class WeightedCentroid(CenterDetectMethodBase):
    def _process(self, *args, **kwargs):
        # 生成横坐标序列
        row_axis = np.ones([self.height, 1], dtype=np.uint32) @ np.arange(
            self.axis_0, self.width + self.axis_0, dtype=np.uint32
        ).reshape([1, self.width])
        # 生成纵坐标序列
        col_axis = np.arange(
            self.axis_0, self.height + self.axis_0, dtype=np.uint32
        ).reshape([self.width, 1]) @ np.ones([1, self.width], dtype=np.uint32)
        # 估计图像中心横坐标：
        square_fig = self.fig.astype(dtype=np.uint32) ** 2
        x_c = np.sum(row_axis * square_fig) / np.sum(square_fig)
        # 估计图像中心纵坐标：
        y_c = np.sum(col_axis * square_fig) / np.sum(square_fig)
        return x_c, y_c

    @property
    def _name(self):
        return "WeightedCentroid"


class HeissanMethod(CenterDetectMethodBase):
    @property
    def _name(self):
        return "HeissanMethod"

    def _process(self, ksize: int, sigma: int, *args, **kwargs):
        gauss_gx = self._gauss_gx(ksize, sigma)
        gauss_gy = self._gauss_gy(ksize, sigma)
        gauss_gxy = self._gauss_gxy(ksize, sigma)
        gauss_gxx = self._gauss_gxx(ksize, sigma)
        gauss_gyy = self._gauss_gyy(ksize, sigma)
        gx = ss.convolve2d(self.fig, gauss_gx, mode="same")
        gy = ss.convolve2d(self.fig, gauss_gy, mode="same")
        gxx = ss.convolve2d(self.fig, gauss_gxx, mode="same")
        gyy = ss.convolve2d(self.fig, gauss_gyy, mode="same")
        gxy = ss.convolve2d(self.fig, gauss_gxy, mode="same")
        C_raw = np.abs(gxx * gyy - gxy**2)
        C_raw = np.abs(gxx + gyy)
        x0_y0 = np.argwhere(C_raw == np.max(C_raw))
        x0, y0 = x0_y0[0, 0], x0_y0[0, 1]
        s = (gy[x0, y0] * gxy[x0, y0] - gx[x0, y0] * gyy[x0, y0]) / (
            gxx[x0, y0] * gyy[x0, y0] - gxy[x0, y0] ** 2
        )
        t = (gx[x0, y0] * gxy[x0, y0] - gy[x0, y0] * gxx[x0, y0]) / (
            gxx[x0, y0] * gyy[x0, y0] - gxy[x0, y0] ** 2
        )
        x0 += s
        y0 += t
        return x0, y0

    # 生成二维高斯x方向一阶导，即gx模板

    def _gauss_gx(self, ksize, sigma):
        kernel = np.zeros((ksize, ksize))
        x_m = np.arange(ksize) - ksize // 2  # [2 1  0  -1  -2]
        y_n = np.arange(ksize) - ksize // 2
        for i in x_m:
            for j in y_n:
                ki = ksize // 2 - i
                kj = ksize // 2 - j
                value = (
                    -i
                    * np.exp(-(i * i + j * j) / (2 * sigma * sigma))
                    / (2 * np.pi * np.power(sigma, 4))
                )
                kernel[kj, ki] = value
        return kernel

    # 生成二维高斯y方向一阶导，即gy模板

    def _gauss_gy(self, ksize, sigma):
        kernel = np.zeros((ksize, ksize))
        x_m = np.arange(ksize) - ksize // 2  # [-2 -1  0  1  2]
        y_n = np.arange(ksize) - ksize // 2
        for i in x_m:
            for j in y_n:
                ki = ksize // 2 - i
                kj = ksize // 2 - j
                value = (
                    -j
                    * np.exp(-(i * i + j * j) / (2 * sigma * sigma))
                    / (2 * np.pi * np.power(sigma, 4))
                )
                kernel[kj, ki] = value
        return kernel

    # 生成二维高斯x方向二阶导，即gxx模板

    def _gauss_gxx(self, ksize, sigma):
        kernel = np.zeros((ksize, ksize))
        x_m = np.arange(ksize) - ksize // 2  # [-2 -1  0  1  2]
        y_n = np.arange(ksize) - ksize // 2
        for i in x_m:
            for j in y_n:
                ki = ksize // 2 - i
                kj = ksize // 2 - j
                value = (
                    (i**2 - sigma**2)
                    * np.exp(-(i * i + j * j) / (2 * sigma * sigma))
                    / (2 * np.pi * np.power(sigma, 6))
                )
                kernel[kj, ki] = value
        return kernel

    # 生成二维高斯y方向二阶导，即gyy模板

    def _gauss_gyy(self, ksize, sigma):
        kernel = np.zeros((ksize, ksize))
        x_m = np.arange(ksize) - ksize // 2  # [-2 -1  0  1  2]
        y_n = np.arange(ksize) - ksize // 2
        for i in x_m:
            for j in y_n:
                ki = ksize // 2 - i
                kj = ksize // 2 - j
                value = (
                    (j**2 - sigma**2)
                    * np.exp(-(i * i + j * j) / (2 * sigma * sigma))
                    / (2 * np.pi * np.power(sigma, 6))
                )
                kernel[kj, ki] = value
        return kernel

    # 生成二维高斯xy方向二阶导，即gxy模板
    def _gauss_gxy(self, ksize, sigma):
        kernel = np.zeros((ksize, ksize))
        x_m = np.arange(ksize) - ksize // 2  # [-2 -1  0  1  2]
        y_n = np.arange(ksize) - ksize // 2
        for i in x_m:
            for j in y_n:
                ki = ksize // 2 - i
                kj = ksize // 2 - j
                value = (
                    i
                    * j
                    * np.exp(-(i * i + j * j) / (2 * sigma * sigma))
                    / (2 * np.pi * np.power(sigma, 6))
                )
                kernel[kj, ki] = value
        return kernel


class GaussFitMethod(CenterDetectMethodBase):
    @property
    def _name(self):
        return "GaussFitMethod"

    def _process(self, *args, **kwargs):
        # 生成横坐标序列
        row_axis = np.ones([self.height, 1], dtype=np.uint32) @ np.arange(
            self.axis_0, self.width + self.axis_0, dtype=np.uint32
        ).reshape([1, self.width])
        # 生成纵坐标序列
        col_axis = np.arange(
            self.axis_0, self.height + self.axis_0, dtype=np.uint32
        ).reshape([self.width, 1]) @ np.ones([1, self.width], dtype=np.uint32)
        X = np.concatenate(
            (
                (row_axis**2).reshape([1, self.height * self.width]),
                (col_axis**2).reshape([1, self.height * self.width]),
                row_axis.reshape([1, self.height * self.width]),
                col_axis.reshape([1, self.height * self.width]),
                np.ones_like(col_axis.reshape([1, self.height * self.width])),
            )
        )
        Y = self.fig.reshape([1, self.height * self.width])
        ind = Y > 0
        Y = Y[ind]
        X = X[:, ind[0]]
        A = (np.log(Y) @ X.T @ np.matrix(X @ X.T).I).tolist()
        A = A[0]
        x_0 = -A[2] / 2 / A[0]
        y_0 = -A[3] / 2 / A[1]
        return x_0, y_0
